#include "ncmcrypt.h"

#include <stdlib.h>
#include <taglib/attachedpictureframe.h>
#include <taglib/flacfile.h>
#include <taglib/id3v2tag.h>
#include <taglib/mpegfile.h>
#include <taglib/tag.h>

#include <cstdint>

//#include <array>
#include <cstdio>
#include <stdexcept>
#include <string>
#include <vector>

#include "aes.h"
#include "base64.h"
#include "cJSON.h"
#include "color.h"

const unsigned char NeteaseCrypt::_sCoreKey[17] = {0x68, 0x7A, 0x48, 0x52, 0x41, 0x6D, 0x73, 0x6F, 0x35,
                                                   0x6B, 0x49, 0x6E, 0x62, 0x61, 0x78, 0x57, 0};
const unsigned char NeteaseCrypt::_sModifyKey[17] = {0x23, 0x31, 0x34, 0x6C, 0x6A, 0x6B, 0x5F, 0x21, 0x5C,
                                                     0x5D, 0x26, 0x30, 0x55, 0x3C, 0x27, 0x28, 0};

const unsigned char NeteaseCrypt::_mPng[8] = {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A};

static void aesEcbDecrypt(const unsigned char* key, std::string& src, std::string& dst)
{
    decltype(src.length()) n = 0;

    unsigned char out[16];

    n = src.length() >> 4;

    dst.clear();

    AES aes(key);

    for (auto i = 0; i < n - 1; i++)
    {
        aes.decrypt((unsigned char*)src.c_str() + (i << 4), out);
        dst += std::string((char*)out, 16);
    }
    auto i = n - 1;
    aes.decrypt((unsigned char*)src.c_str() + (i << 4), out);
    auto pad = out[15];
    if (pad > 16)
    {
        pad = 0;
    }
    dst += std::string((char*)out, 16 - pad);
}

static void Replace(std::string& str, const std::string& from, const std::string& to)
{
    if (from.empty()) return;
    size_t start_pos = 0;
    while ((start_pos = str.find(from, start_pos)) != std::string::npos)
    {
        str.replace(start_pos, from.length(), to);
        start_pos += to.length();  // In case 'to' contains 'from', like replacing 'x' with 'yx'
    }
}

static std::string FileNameWithoutExt(const std::string& str)
{
    size_t lastPath = str.find_last_of("/\\"); /* 寻找 / or \\ */
    std::string path = str.substr(lastPath + 1);
    size_t lastExt = path.find_last_of(".");
    return path.substr(0, lastExt);
}

NeteaseMusicMetadata::~NeteaseMusicMetadata() { cJSON_Delete(mRaw); }

NeteaseMusicMetadata::NeteaseMusicMetadata(cJSON* raw)
{
    if (!raw)
    {
        return;
    }

    cJSON* iSwap;
    int artistLen, i;

    mRaw = raw;

    iSwap = cJSON_GetObjectItem(raw, "musicName");
    if (iSwap)
    {
        mName = std::string(cJSON_GetStringValue(iSwap));
    }

    iSwap = cJSON_GetObjectItem(raw, "album");
    if (iSwap)
    {
        mAlbum = std::string(cJSON_GetStringValue(iSwap));
    }

    iSwap = cJSON_GetObjectItem(raw, "artist");
    if (iSwap)
    {
        artistLen = cJSON_GetArraySize(iSwap);

        for (int i = 0; i < artistLen - 1; i++)
        {
            mArtist += std::string(cJSON_GetStringValue(cJSON_GetArrayItem(cJSON_GetArrayItem(iSwap, i), 0)));
            mArtist += "/";
        }
        mArtist += std::string(cJSON_GetStringValue(cJSON_GetArrayItem(cJSON_GetArrayItem(iSwap, artistLen - 1), 0)));
    }

    iSwap = cJSON_GetObjectItem(raw, "bitrate");
    if (iSwap)
    {
        mBitrate = iSwap->valueint;
    }

    iSwap = cJSON_GetObjectItem(raw, "duration");
    if (iSwap)
    {
        mDuration = iSwap->valueint;
    }

    iSwap = cJSON_GetObjectItem(raw, "format");
    if (iSwap)
    {
        mFormat = std::string(cJSON_GetStringValue(iSwap));
    }
}

bool NeteaseCrypt::OpenFile(std::string const& path)
{
    try
    {
        _mFile.open(path, std::ios::in | std::ios::binary);
    }
    catch (...)
    {
        return false;
    }
    return true;
}

/*!
 * @brief 通过头几个 byte，检测是否为 ncm 文件
 * @return
*/
bool NeteaseCrypt::IsNcmFile()
{
    // ncm MAGIC_HEADER	8 bytes; ncm 文件的文件头标识;  内容为 b"CTENFDAM"
    // bCTENFDAM: 43 54 45 4E 46 44 41 4D
    // constexpr auto size = 4;  // 每次读取 4 byte,
    uint32_t header = 0;
    _mFile.read(reinterpret_cast<char*>(&header), sizeof(header));
    // 前四位是 43 54 45 4E
    // 但是按照 little endian, 高内存在高位, 所以是反向的
    if (header != (unsigned int)0x4e455443)
    {
        return false;
    }
    // 前四位是 46 44 41 4D, 类似地，反向
    _mFile.read(reinterpret_cast<char*>(&header), sizeof(header));
    if (header != (unsigned int)0x4d414446)
    {
        return false;
    }

    return true;
}

/*!
 * @brief 返回读取的 byte 数
 * @param s
 * @param n
 * @return
*/
std::streamsize NeteaseCrypt::Read(std::istream::char_type* s, std::streamsize n)
{
    _mFile.read(s, n);

    auto gcount = _mFile.gcount();

    if (gcount <= 0)
    {
        throw std::invalid_argument("can't read file");
    }

    return gcount;
}

void NeteaseCrypt::BuildKeyBox(unsigned char* key, int keyLen)
{
    for (int i = 0; i < 256; ++i)
    {
        _mKeyBox[i] = (unsigned char)i;
    }

    unsigned char swap = 0;
    unsigned char c = 0;
    unsigned char last_byte = 0;
    unsigned char key_offset = 0;

    for (int i = 0; i < 256; ++i)
    {
        swap = _mKeyBox[i];
        c = ((swap + last_byte + key[key_offset++]) & 0xff);
        if (key_offset >= keyLen) key_offset = 0;
        _mKeyBox[i] = _mKeyBox[c];
        _mKeyBox[c] = swap;
        last_byte = c;
    }
}

std::string NeteaseCrypt::MimeType(const char* data)
{
    using namespace std::literals;
    if (memcmp(data, _mPng, 8) == 0)
    {
        return "image/png"s;
    }

    return "image/jpeg"s;
}

void NeteaseCrypt::FixMetadata()
{
    if (_mDumpFilepath.length() <= 0)
    {
        throw std::invalid_argument("must dump before");
    }

    TagLib::File* audioFile{};
    TagLib::Tag* tag{};
    TagLib::ByteVector ivec(_mImageData, _mImageLen);

    if (_mFormat == NeteaseCrypt::MP3)
    {
        audioFile = new TagLib::MPEG::File(_mDumpFilepath.c_str());
        tag = dynamic_cast<TagLib::MPEG::File*>(audioFile)->ID3v2Tag(true);

        if (_mImageLen > 0)
        {
            auto* const frame = new TagLib::ID3v2::AttachedPictureFrame;

            const auto& itype = TagLib::String(_mImageData, TagLib::String::UTF8);
            frame->setMimeType(itype);
            frame->setPicture(ivec);

            dynamic_cast<TagLib::ID3v2::Tag*>(tag)->addFrame(frame);
        }
    }
    else if (_mFormat == NeteaseCrypt::FLAC)
    {
        audioFile = new TagLib::FLAC::File(_mDumpFilepath.c_str());
        tag = audioFile->tag();

        if (_mImageLen > 0)
        {
            TagLib::FLAC::Picture* cover = new TagLib::FLAC::Picture;
            MimeType(_mImageData);
            cover->setMimeType(MimeType(_mImageData));
            cover->setType(TagLib::FLAC::Picture::FrontCover);
            cover->setData(ivec);

            dynamic_cast<TagLib::FLAC::File*>(audioFile)->addPicture(cover);
        }
    }

    if (_mMetaData != NULL)
    {
        tag->setTitle(TagLib::String(_mMetaData->Name(), TagLib::String::UTF8));
        tag->setArtist(TagLib::String(_mMetaData->Artist(), TagLib::String::UTF8));
        tag->setAlbum(TagLib::String(_mMetaData->Album(), TagLib::String::UTF8));
    }

    tag->setComment(TagLib::String("Create by netease copyright protected dump tool. author 5L", TagLib::String::UTF8));

    audioFile->save();
}

void NeteaseCrypt::Dump()
{
    _mDumpFilepath.clear();
    _mDumpFilepath.resize(1024);

    if (_mMetaData)
    {
        _mDumpFilepath = _mMetaData->Name();

        //Replace(mDumpFilepath_, "\\", "＼");
        Replace(_mDumpFilepath, "/", "_");
        Replace(_mDumpFilepath, "?", "_");
        Replace(_mDumpFilepath, ":", "_");
        Replace(_mDumpFilepath, "*", "_");
        Replace(_mDumpFilepath, "\"", "_");
        Replace(_mDumpFilepath, "<", "_");
        Replace(_mDumpFilepath, ">", "_");
        Replace(_mDumpFilepath, "|", "_");
    }
    //else
    //{
    _mDumpFilepath = FileNameWithoutExt(_mFilepath);  // file dir + file baseName
    //}

    constexpr std::streamsize mySize = 0x8000;
    std::vector<std::istream::char_type> buffer(mySize);
    //buffer.resize(size);
    //std::array<unsigned char, size> buffer{0};
    //unsigned char buffer[size] = { 0 };
    std::ofstream output;

    while (!_mFile.eof())
    {
        const auto n = Read(buffer.data(), mySize);

        // RC4算法的第二部分，伪随机序列产生算法（Pseudo Random Generation Algorithm，PRGA）
        for (auto i = 0; i < n; i++)
        {
            int j = (i + 1) & 0xff;
            buffer[i] ^= _mKeyBox[(_mKeyBox[j] + _mKeyBox[(_mKeyBox[j] + j) & 0xff]) & 0xff];
        }

        if (!output.is_open())
        {
            // identify format; ID3 format mp3
            // 49       44      33      例如 04
            // I(73)    D(68)   3(51)   例如 4(V2.4)
            if (buffer[0] == 0x49 && buffer[1] == 0x44 && buffer[2] == 0x33)
            {
                _mDumpFilepath += ".mp3";
                _mFormat = NeteaseCrypt::NcmFormat::MP3;
            }
            else
            {
                // FLAC 的文件头由4个字节的 ASCII 码组成的标志"fLaC"，16进为 66 4C 61 43
                _mDumpFilepath += ".flac";
                _mFormat = NeteaseCrypt::NcmFormat::FLAC;
            }

            output.open(_mDumpFilepath, output.out | output.binary);
        }

        output.write((char*)buffer.data(), n);
    }

    output.flush();
    output.close();
}

NeteaseCrypt::~NeteaseCrypt()
{
    if (_mMetaData != NULL)
    {
        delete _mMetaData;
    }

    _mFile.close();
}

/*!
 * @brief 读取 解密 meta data
 * @param path
*/
NeteaseCrypt::NeteaseCrypt(std::string const& path)
{
    if (!OpenFile(path))
    {
        throw std::invalid_argument("can't open file");
    }

    if (!IsNcmFile())
    {
        throw std::invalid_argument("not netease protected file");
    }

    // 跳过 b"\x01\x61"; 新版是 x6D, x01
    std::istream::off_type offset = 2;
    if (!_mFile.seekg(offset, _mFile.cur))
    {
        throw std::invalid_argument("can't seek file");
    }

    _mFilepath = path;

    size_t rc4_key_enc_size = 0;  // rc4_key_enc_size，加密的 RC4算法密钥 的长度,
    constexpr size_t size = 4;    // 固定长度, 每次读取量
    this->Read(reinterpret_cast<char*>(&rc4_key_enc_size), size);
    //endian_swap(fileSize);
    if (rc4_key_enc_size <= 0)
    {
        throw std::invalid_argument("broken ncm file");
    }
    // read rc4_key_enc
    std::vector<char> keydata(rc4_key_enc_size);
    this->Read(keydata.data(), rc4_key_enc_size);

    for (int i = 0; i < rc4_key_enc_size; i++)
    {
        keydata[i] ^= 0x64;
    }

    std::string rawKeyData(keydata.data(), rc4_key_enc_size);
    std::string mKeyData;

    aesEcbDecrypt(_sCoreKey, rawKeyData, mKeyData);

    BuildKeyBox((unsigned char*)mKeyData.c_str() + 17, mKeyData.length() - 17);

    size_t metadata_enc_size = 0;  //metadata_enc_size 4 byte int
    Read(reinterpret_cast<char*>(&metadata_enc_size), size);

    if (metadata_enc_size <= 0)
    {
        printf("[Warn] `%s` missing metadata infomation can't fix some infomation!\n", path.c_str());

        _mMetaData = NULL;
    }
    else
    {
        std::vector<char> modifyData(metadata_enc_size);
        Read(modifyData.data(), metadata_enc_size);

        for (auto i = 0; i < metadata_enc_size; i++)
        {
            // 0110 0011 混淆加密
            modifyData[i] ^= 0x63;
        }

        std::string swapModifyData;
        std::string modifyOutData;
        std::string modifyDecryptData;  // 歌曲的 meta data; json 格式

        swapModifyData = std::string(modifyData.data() + 22, metadata_enc_size - 22);

        // 移除 `163 key(Don't modify):`
        Base64::Decode(swapModifyData, modifyOutData);

        aesEcbDecrypt(_sModifyKey, modifyOutData, modifyDecryptData);

        // 移除 `music:`
        modifyDecryptData = std::string(modifyDecryptData.begin() + 6, modifyDecryptData.end());

        _mMetaData = new NeteaseMusicMetadata(cJSON_Parse(modifyDecryptData.c_str()));
    }

    //********************************* skip crc32 & unuse charset
    // skip crc32 & image version
    if (!_mFile.seekg(5, _mFile.cur))
    {
        throw std::invalid_argument("can't seek file");
    }

    unsigned int cover_size;
    uint32_t cover_frame_len{0};
    Read(reinterpret_cast<char*>(&cover_frame_len), 4);
    Read(reinterpret_cast<char*>(&cover_size), sizeof(cover_size));

    if (cover_size > 0)
    {
        _mImageData = std::string(cover_size, '\0').data();
        Read(&_mImageData[0], cover_size);
    }
    else
    {
        std::cout << BOLDYELLOW << "[Warn] " << RESET << "'" << path << "' missing album can't fix album image!"
                  << std::endl;
    }
    _mFile.seekg(cover_frame_len - cover_size, _mFile.cur);
}

// 交换 little endian, big endian 顺序
void NeteaseCrypt::EndianSwap(unsigned int& x)
{
    x = (x >> 24) | ((x << 8) & 0x00FF0000) | ((x >> 8) & 0x0000FF00) | (x << 24);
}
